#!/usr/bin/env python3
"""
简洁SD3 LoRA训练参数节点
只保留核心必需的参数，避免nan损失
包含缓存参数节点和参数合并器节点
"""

import folder_paths
import os
from .lora_trainer_utils.constants import SD3_TRAINING_PARAMS_CATEGORY, SD3_TRAINING_PARAMS, SD3_CACHE_PARAMS, SD3_SAMPLING_PARAMS, SD3_BUCKET_PARAMS

class SD3TrainingParams:
    """简洁SD3 LoRA训练参数节点"""
    
    aux_id = "comfyui_lora_train"
    
    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                # 基础模型与数据
                "pretrained_model_name_or_path": (folder_paths.get_filename_list("checkpoints"), {
                    "default": "",
                    "tooltip": "预训练模型路径（SD3/SD3.5底模）"
                }),
                "train_data_dir": ("STRING", {
                    "default": "",
                    "tooltip": "训练图片文件夹"
                }),
                "output_dir": ("STRING", {
                    "default": "./output",
                    "tooltip": "模型输出目录"
                }),
                "output_name": ("STRING", {
                    "default": "sd3_lora",
                    "tooltip": "输出模型名称"
                }),
                
                # LoRA结构 - 简化
                "network_dim": ("INT", {
                    "default": 32,
                    "min": 1,
                    "max": 64,
                    "tooltip": "LoRA网络维度（RANK），推荐16-32，过大可能导致过拟合"
                }),
                "network_alpha": ("INT", {
                    "default": 16,
                    "min": 1,
                    "max": 64,
                    "tooltip": "LoRA alpha，建议为dim的1/2"
                }),
                
                # 训练核心参数 - 简化
                "learning_rate": ("FLOAT", {
                    "default": 5e-5,
                    "min": 1e-4,
                    "max": 1e-3,
                    "tooltip": "学习率，推荐5e-5，过高可能导致损失过大"
                }),
                "max_train_epochs": ("INT", {
                    "default": 3,
                    "min": 1,
                    "max": 100,
                    "tooltip": "最大训练轮数，建议3-6轮"
                }),
                "train_batch_size": ("INT", {
                    "default": 2,
                    "min": 1,
                    "max": 50,
                    "tooltip": "训练批次大小"
                }),
                "resolution": ("INT", {
                    "default": 1024,
                    "min": 512,
                    "max": 1024,
                    "tooltip": "训练分辨率，推荐768"
                }),
                
                # 优化器 - 简化
                "optimizer_type": (["AdamW8bit", "AdamW", "AdaFactor"], {
                    "default": "AdamW",
                    "tooltip": "优化器类型，推荐AdamW8bit，更稳定"
                }),
                "lr_scheduler": (["cosine", "linear", "constant"], {
                    "default": "linear",
                    "tooltip": "学习率调度器，推荐linear，更稳定"
                }),
                
                # 混合精度 - 强制bf16
                "mixed_precision": (["bf16"], {
                    "default": "bf16",
                    "tooltip": "混合精度训练，强制bf16防止nan"
                }),
                
                # 保存设置
                "save_every_n_steps": ("INT", {
                    "default": 500,
                    "min": 100,
                    "max": 5000,
                    "tooltip": "每多少步保存一次模型"
                }),
                "save_model_as": (["safetensors"], {
                    "default": "safetensors",
                    "tooltip": "模型保存格式"
                }),
            },
            "optional": {
                # SD3特有模型路径
                "clip_l": (["None"] + folder_paths.get_filename_list("clip"), {
                    "default": "None",
                    "tooltip": "CLIP-L模型路径，选择None时使用基础模型内置CLIP-L"
                }),
                "clip_g": (folder_paths.get_filename_list("clip"), {
                    "default": "",
                    "tooltip": "CLIP-G模型路径，SD3专用，必需参数"
                }),
                "t5xxl": (["None"] + folder_paths.get_filename_list("text_encoders"), {
                    "default": "None",
                    "tooltip": "T5-XXL模型路径，选择None时使用基础模型内置T5-XXL"
                }),
                "vae": (["None"] + folder_paths.get_filename_list("vae"), {
                    "default": "None",
                    "tooltip": "可选VAE模型路径，选择None时使用基础模型内置VAE"
                }),
                
                # 文本编码器训练 - 简化
                "train_text_encoder": ("BOOLEAN", {
                    "default": False,
                    "tooltip": "是否训练文本编码器（CLIP-L和CLIP-G）"
                }),
                "train_t5xxl": ("BOOLEAN", {
                    "default": False,
                    "tooltip": "是否训练T5-XXL"
                }),
                
                # 性能优化 - 简化
                "gradient_checkpointing": ("BOOLEAN", {
                    "default": True,
                    "tooltip": "梯度检查点，节省显存"
                }),
                "max_data_loader_n_workers": ("INT", {
                    "default": 2,
                    "min": 0,
                    "max": 8,
                    "tooltip": "数据加载器工作进程数"
                }),
                "persistent_data_loader_workers": ("BOOLEAN", {
                    "default": True,
                    "tooltip": "持久化数据加载器工作进程"
                }),
                
                # 随机种子
                "seed": ("INT", {
                    "default": 42,
                    "min": 0,
                    "max": 2147483647,
                    "tooltip": "随机种子"
                }),
                
                # 梯度裁剪
                "max_grad_norm": ("FLOAT", {
                    "default": 1.0,
                    "min": 0.0,
                    "max": 1.0,
                    "tooltip": "梯度裁剪阈值，防止梯度爆炸"
                }),
                
                # 数据增强设置 - sd-scripts原生支持
                "random_crop": ("BOOLEAN", {
                    "default": True,
                    "tooltip": "启用随机裁剪，解决图片尺寸过大问题（与缓存冲突）"
                }),
                "flip_aug": ("BOOLEAN", {
                    "default": True,
                    "tooltip": "启用水平翻转增强，提升数据多样性"
                }),
                "color_aug": ("BOOLEAN", {
                    "default": False,
                    "tooltip": "启用颜色增强，提升模型泛化能力（与缓存冲突）"
                }),
                
                # SD3特有优化 - 简化
                "fp8_base": ("BOOLEAN", {
                    "default": True,
                    "tooltip": "启用fp8基础模型，节省显存"
                }),
                "blocks_to_swap": ("INT", {
                    "default": 22,
                    "min": 0,
                    "max": 36,
                    "tooltip": "内存优化：交换块数，SD3.5L最多36块，SD3.5M最多22块"
                }),
                "disable_mmap_load_safetensors": ("BOOLEAN", {
                    "default": True,
                    "tooltip": "禁用safetensors内存映射加载，节省显存"
                }),
                

                
                # 网络训练设置
                "network_train_unet_only": ("BOOLEAN", {
                    "default": True,
                    "tooltip": "仅训练UNet，不训练文本编码器"
                }),
            }
        }
    
    @classmethod
    def VALIDATE_INPUTS(cls, **kwargs):
        """验证输入参数"""
        # 检查训练数据目录是否存在
        train_data_dir = kwargs.get("train_data_dir")
        if train_data_dir and isinstance(train_data_dir, str) and train_data_dir.strip() and not os.path.exists(train_data_dir):
            return f"训练数据目录不存在: {train_data_dir}"
        
        return True
    
    RETURN_TYPES = (SD3_TRAINING_PARAMS,)
    RETURN_NAMES = ("训练参数",)
    FUNCTION = "generate_params"
    CATEGORY = SD3_TRAINING_PARAMS_CATEGORY
    
    def generate_params(self, **kwargs):
        """生成训练参数 - 简化版本"""
        
        # 处理"None"值，删除对应参数
        none_params = ["vae", "clip_l", "t5xxl"]
        for param in none_params:
            if kwargs.get(param) == "None":
                kwargs.pop(param, None)
                print(f"[INFO] 参数 {param} 设置为None，已从训练参数中移除")
        
        # 自动处理所有常见路径参数，转为绝对路径
        path_keys = [
            ("pretrained_model_name_or_path", "checkpoints"),
            ("vae", "vae"),
            ("clip_l", "clip"),
            ("clip_g", "clip"),
            ("t5xxl", "text_encoders"),
        ]
        for key, folder in path_keys:
            if kwargs.get(key):
                try:
                    if folder:
                        kwargs[key] = folder_paths.get_full_path(folder, kwargs[key])
                    else:
                        if not os.path.isabs(kwargs[key]):
                            kwargs[key] = os.path.abspath(kwargs[key])
                except Exception as e:
                    print(f"[WARNING] 获取{key}完整路径失败: {e}")
        
        # 处理分辨率格式：转为"width,height"格式
        if "resolution" in kwargs:
            resolution = kwargs["resolution"]
            kwargs["resolution"] = f"{resolution},{resolution}"
        
        # 处理输出路径：检查并创建目录
        if kwargs.get("output_dir"):
            output_dir = kwargs["output_dir"]
            if not os.path.isabs(output_dir):
                output_dir = os.path.abspath(output_dir)
            kwargs["output_dir"] = output_dir
            
            # 创建输出目录（如果不存在）
            try:
                os.makedirs(output_dir, exist_ok=True)
                print(f"[INFO] 输出目录已创建/确认: {output_dir}")
            except Exception as e:
                print(f"[WARNING] 创建输出目录失败: {e}")
        
        # 处理训练数据目录：检查是否存在
        if kwargs.get("train_data_dir"):
            train_data_dir = kwargs["train_data_dir"]
            if not os.path.isabs(train_data_dir):
                train_data_dir = os.path.abspath(train_data_dir)
            kwargs["train_data_dir"] = train_data_dir
            
            if not os.path.exists(train_data_dir):
                print(f"[WARNING] 训练数据目录不存在: {train_data_dir}")
        

        kwargs["network_module"] = "networks.lora_sd3"
        
        # 处理训练模式互斥
        if kwargs.get("network_train_unet_only", False) and kwargs.get("network_train_text_encoder_only", False):
            print("[WARNING] 不能同时仅训练UNet和仅训练文本编码器，将使用UNet训练模式")
            kwargs["network_train_text_encoder_only"] = False
        
        print("[INFO] 已应用简化SD3训练参数配置")
        
        return (kwargs,)
    



class SD3CacheParams:
    """SD3缓存参数节点"""
    
    aux_id = "comfyui_lora_train"
    
    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                # 文本编码器缓存
                "cache_text_encoder_outputs": ("BOOLEAN", {
                    "default": True,
                    "tooltip": "缓存文本编码器输出，节省显存和加速训练"
                }),
                "cache_text_encoder_outputs_to_disk": ("BOOLEAN", {
                    "default": False,
                    "tooltip": "将文本编码器输出缓存到磁盘，节省显存但可能稍慢"
                }),
                "use_t5xxl_cache_only": ("BOOLEAN", {
                    "default": False,
                    "tooltip": "仅缓存T5-XXL输出，不缓存CLIP输出"
                }),
                
                # Latent缓存
                "cache_latents": ("BOOLEAN", {
                    "default": True,
                    "tooltip": "缓存Latents，加速训练"
                }),
                "cache_latents_to_disk": ("BOOLEAN", {
                    "default": False,
                    "tooltip": "将latents缓存到磁盘，节省显存"
                }),
                "vae_batch_size": ("INT", {
                    "default": 1,
                    "min": 1,
                    "max": 8,
                    "tooltip": "VAE批处理大小，用于latent缓存"
                }),
                
                # 其他缓存设置
                "skip_cache_check": ("BOOLEAN", {
                    "default": False,
                    "tooltip": "跳过缓存检查，直接使用现有缓存"
                }),
                "tokenizer_cache_dir": ("STRING", {
                    "default": "",
                    "tooltip": "Tokenizer缓存目录，留空使用默认目录"
                }),
                
                # 数据加载器缓存
                "max_data_loader_n_workers": ("INT", {
                    "default": 2,
                    "min": 0,
                    "max": 8,
                    "tooltip": "数据加载器工作进程数"
                }),
                "persistent_data_loader_workers": ("BOOLEAN", {
                    "default": True,
                    "tooltip": "持久化数据加载器工作进程"
                }),
            },
            "optional": {
                # 高级缓存设置
            }
        }
    
    @classmethod
    def VALIDATE_INPUTS(cls, **kwargs):
        """验证输入参数"""
        # 检查tokenizer缓存目录
        tokenizer_cache_dir = kwargs.get("tokenizer_cache_dir")
        if tokenizer_cache_dir and isinstance(tokenizer_cache_dir, str) and tokenizer_cache_dir.strip():
            if not os.path.exists(tokenizer_cache_dir):
                try:
                    os.makedirs(tokenizer_cache_dir, exist_ok=True)
                except Exception as e:
                    return f"无法创建tokenizer缓存目录: {e}"
        
        return True
    
    RETURN_TYPES = ("SD3_CACHE_PARAMS",)
    RETURN_NAMES = ("缓存参数",)
    FUNCTION = "generate_cache_params"
    CATEGORY = SD3_TRAINING_PARAMS_CATEGORY
    
    def generate_cache_params(self, **kwargs):
        """生成缓存参数"""
        
        
        # 处理tokenizer缓存目录
        tokenizer_cache_dir = kwargs.get("tokenizer_cache_dir", "")
        if tokenizer_cache_dir :
            try:
                if not os.path.isabs(tokenizer_cache_dir):
                    os.makedirs(tokenizer_cache_dir, exist_ok=True)
                    print(f"[INFO] 创建tokenizer缓存目录: {tokenizer_cache_dir}")
            except Exception as e:
                print(f"[WARNING] 创建tokenizer缓存目录失败: {e}")
                kwargs.pop("tokenizer_cache_dir", None)
        else:
            kwargs.pop("tokenizer_cache_dir", None)
        
        return (kwargs,)


class SD3BucketParams:
    """SD3分桶参数节点"""
    
    aux_id = "comfyui_lora_train"
    
    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                # 分桶设置
                "enable_bucket": ("BOOLEAN", {
                    "default": True,
                    "tooltip": "启用图像分桶，提高训练效率"
                }),
                "min_bucket_reso": ("INT", {
                    "default": 256,
                    "min": 64,
                    "max": 1024,
                    "tooltip": "分桶最小分辨率"
                }),
                "max_bucket_reso": ("INT", {
                    "default": 1024,
                    "min": 64,
                    "max": 2048,
                    "tooltip": "分桶最大分辨率"
                }),
                "bucket_reso_steps": ("INT", {
                    "default": 64,
                    "min": 1,
                    "max": 512,
                    "tooltip": "分桶分辨率步长"
                }),
            },
            "optional": {
                # 分桶优化
                "highvram": ("BOOLEAN", {
                    "default": False,
                    "tooltip": "高显存模式，显存充足时启用"
                }),
                "sdpa": ("BOOLEAN", {
                    "default": True,
                    "tooltip": "启用Scaled Dot-Product Attention，加速训练"
                }),
            }
        }
    
    @classmethod
    def VALIDATE_INPUTS(cls, **kwargs):
        """验证输入参数"""
        # 检查分桶最小分辨率是否小于最大分辨率
        min_bucket_reso = kwargs.get("min_bucket_reso")
        max_bucket_reso = kwargs.get("max_bucket_reso")
        bucket_reso_steps = kwargs.get("bucket_reso_steps")
        
        if min_bucket_reso is not None and max_bucket_reso is not None and bucket_reso_steps is not None:
            if min_bucket_reso >= max_bucket_reso:
                return "分桶最小分辨率必须小于最大分辨率"
            if bucket_reso_steps <= 0:
                return "分桶分辨率步长必须大于0"
        
        return True
    
    RETURN_TYPES = ("SD3_BUCKET_PARAMS",)
    RETURN_NAMES = ("分桶参数",)
    FUNCTION = "generate_bucket_params"
    CATEGORY = SD3_TRAINING_PARAMS_CATEGORY
    
    def generate_bucket_params(self, **kwargs):
        """生成分桶参数"""
        return (kwargs,)


class SD3ParamsMerger:
    """SD3参数合并节点"""
    
    aux_id = "comfyui_lora_train"
    
    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "training_params": (SD3_TRAINING_PARAMS, {
                    "tooltip": "SD3训练参数"
                }),
            },
            "optional": {
                "cache_params": (SD3_CACHE_PARAMS, {
                    "tooltip": "SD3缓存参数（可选）"
                }),
                "bucket_params": ("SD3_BUCKET_PARAMS", {
                    "tooltip": "SD3分桶参数（可选）"
                }),
                "sampling_params": (SD3_SAMPLING_PARAMS, {
                    "tooltip": "SD3采样参数（可选）"
                }),
            }
        }
    
    RETURN_TYPES = (SD3_TRAINING_PARAMS,)
    RETURN_NAMES = ("完整训练参数",)
    FUNCTION = "merge_params"
    CATEGORY = SD3_TRAINING_PARAMS_CATEGORY
    
    def merge_params(self, training_params, cache_params=None, bucket_params=None, sampling_params=None):
        """合并训练参数、缓存参数和采样参数"""
        
        # 创建合并后的参数字典
        merged_params = training_params.copy()
        
        # 合并缓存参数（如果提供）
        if cache_params:
            for key, value in cache_params.items():
                merged_params[key] = value
        
        # 合并分桶参数（如果提供）
        if bucket_params:
            for key, value in bucket_params.items():
                merged_params[key] = value
        
        # 合并采样参数（如果提供）
        if sampling_params:
            for key, value in sampling_params.items():
                merged_params[key] = value
        
        # 缓存优先：当启用缓存时自动禁用冲突的数据增强
        if merged_params.get("cache_latents", False):
            if merged_params.get("color_aug", False):
                merged_params["color_aug"] = False
                print("[INFO] 检测到cache_latents启用，自动禁用color_aug以避免冲突")
            
            if merged_params.get("random_crop", False):
                merged_params["random_crop"] = False
                print("[INFO] 检测到cache_latents启用，自动禁用random_crop以避免冲突")
            
            if merged_params.get("flip_aug", False):
                merged_params["flip_aug"] = False
                print("[INFO] 检测到cache_latents启用，自动禁用flip_aug以避免冲突")
        
        return (merged_params,)


# 注册节点
NODE_CLASS_MAPPINGS = {
    "SD3TrainingParams": SD3TrainingParams,
    "SD3CacheParams": SD3CacheParams,
    "SD3BucketParams": SD3BucketParams,
    "SD3ParamsMerger": SD3ParamsMerger,
}

NODE_DISPLAY_NAME_MAPPINGS = {
    "SD3TrainingParams": "🛠️ SD3训练参数(简化版)",
    "SD3CacheParams": "💾 SD3缓存参数",
    "SD3BucketParams": "🪣 SD3分桶参数",
    "SD3ParamsMerger": "🔄 SD3参数合并器",
} 